# coding=utf-8

import random
import numpy as np
import collections

data_file = "Data/poetry.txt"

class Data:
    def __init__(self,filename=None):
        self.filename = filename==None and filename or data_file

        self.poetrys = self.ReadFromDataFile()

        self.words_num_map,self.poetrys_vectors = self.ExtractWordsFromPoetry()
        self.words_length = len(self.words_num_map)
        self.poetrys_num = len(self.poetrys_vectors)
        #self.poetry_length = len(self.poetrys_vectors[0])
        #self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y = self.DataShuffle(self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y)

        self.index = 0
    
    def GetNext_(self,batch_size=64):
        x,y,overflow = self.GetNext(batch_size)
        for s in range(len(x)):
            print "\n\nindex:",s,"\n"
            print "x:"
            for i in range(len(x[0])):
                print self.GetWordFromNum(x[s][i]),
            print ""
            print "y:"
            for i in range(len(y[0])):
                print self.GetWordFromNum(y[s][i]),
            print ""
        return x,y,overflow

    def GetNext(self,batch_size=64):
        start = self.index
        self.index += batch_size
        epoch_overflow = False
        if self.index >= self.poetrys_num: 
            #self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y = self.DataShuffle(self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y)
            start = 0
            self.index = batch_size
            epoch_overflow = True


        self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y = self.ReshapeForRNN(self.poetrys_vectors[start:self.index],batch_size)
        return self.poetrys_vectors_for_rnn_x,self.poetrys_vectors_for_rnn_y,epoch_overflow

    def DataShuffle(self,*args):
        return args
        tmp = zip(*args)
        random.shuffle(tmp)
        tmp_ = []  
        for i in range(len(args)):
            tmp_.append([ one[i] for one in tmp ])
        return tmp_

    def GetWordFromNum(self,num):
        for word in self.words_num_map:
            if self.words_num_map.get(word) == num:
                return word 

    def GetNumFromWord(self,word):
        return self.words_num_map.get(word)

    def GetWordsLen(self):
        return self.words_length

    def GetPoetryLen(self):
        return self.poetry_length

    def ReshapeForRNN(self,data,batch_size=64):
        ''' in this function, firstly, we add all poetrys to the same length with whitespace. Secondly, we move a copy of the data a step towards to the end. And the copy data is designed to be the output for rnn'''
        length = max(map(len,data))
        x_data = np.full([batch_size,length],self.GetNumFromWord(' '),np.int32)
        for row in range(len(data)):
            x_data[row,:len(data[row])] = data[row]

        y_data = np.copy(x_data)
        y_data[:,:-1] = x_data[:,1:]
        return x_data,y_data
            


    def ExtractWordsFromPoetry(self):
        all_words = []
        for poetry in self.poetrys:
            all_words += [word for word in poetry]
        counter = collections.Counter(all_words)
        count_pairs = sorted(counter.items(),key=lambda x:-x[1]) # sorted from less times to high times
        words,_=zip(*count_pairs) # use collections to sorted the words

        words = words[:len(words)] + (' ',) # add a " " to the words list
        words_num_map = dict(zip(words,range(len(words)))) #connect the word with the vector
        
        to_num = lambda word: words_num_map.get(word, len(words))
        poetrys_vectors = [list(map(to_num,poetry)) for poetry in self.poetrys] # convert the poetry to vector list
        return words_num_map,poetrys_vectors

    def ReadFromDataFile(self):
        ''' get a poetrys list sorted by the length'''
        f = open(self.filename,"r")
        poetrys = []
        for line in f.readlines():
            try:
                title,content = line.strip().split(":")
                content = content.replace(" ",",")
                if '_' in content or '(' in content or '（' in content or '《' in content or '【' in content:
                    continue
                if len(content) < 5 or len(content) > 79: continue
                content = "[%s]" % content
                poetrys.append(content.decode("utf-8"))
            except Exception as e:
                pass
        f.close()
        poetrys = sorted(poetrys,key=lambda line:len(line))
        return poetrys

if __name__ == "__main__":
    data = Data()
    print data.GetNext_(128)
    #print "words",data.words_num_map
